# 运行前pip安装vanna库：pip install 'vanna[mysql]'
import vanna
from vanna.remote import VannaDefault
#构建vanna对象
api_key = vanna.get_api_key('2871135045@qq.com') # Your API key from https://vanna.ai/account/profile 
vanna_model_name = 'guapishengteng'# Your model name from https://vanna.ai/account/profile 
vn = VannaDefault(model=vanna_model_name, api_key=api_key)

#链接数据库
con = vn.connect_to_mysql(host='192.168.0.205',dbname = 'vannadb',user = 'root',password='Sp091212',port = 3306)

#添加训练数据
# SQL 语句：创建表
create_table_sql = """
CREATE TABLE IF NOT EXISTS vuser (
    `id` INT PRIMARY KEY COMMENT '用户ID', 
    username VARCHAR(50) COMMENT '用户名',
    email VARCHAR(100) COMMENT '电子邮件', 
    age INT COMMENT '年龄',
    gender VARCHAR(10) COMMENT '性别（男/女）',
    city VARCHAR(50) COMMENT '城市'
) COMMENT='用户信息表' CHARACTER SET=utf8 COLLATE=utf8_unicode_ci;
"""
# 执行创建表的操作
try:
    with con.cursor() as cursor:
        cursor.execute(create_table_sql)
        con.commit()
        print("表创建成功！")
except pymysql.MySQLError as e:
    print(f"创建表失败: {e}")

# SQL 语句：插入数据
insert_data_sql = """
INSERT INTO vuser (`id`, username, email, age, gender, city) VALUES
(1, '张三', 'zhangsan@example.com', 30, '男', '北京'),
(2, '李四', 'lisi@example.com', 25, '女', '上海'),
(3, '王五', 'wangwu@example.com', 40, '男', '广州'),
(4, '赵六', 'zhaoliu@example.com', 35, '女', '深圳'),
(5, '小明', 'xiaoming@example.com', 28, '男', '成都'),
(6, '小红', 'xiaohong@example.com', 45, '女', '重庆'),
(7, '小华', 'xiaohua@example.com', 32, '男', '天津'),
(8, '小丽', 'xiaoli@example.com', 27, '女', '南京'),
(9, '小李', 'xiaoli2@example.com', 38, '男', '武汉'),
(10, '小美', 'xiaomei@example.com', 33, '女', '西安');
"""
# 执行插入数据的操作
try:
    with connection.cursor() as cursor:
        cursor.execute(insert_data_sql)
        connection.commit()
        print("数据插入成功！")
except pymysql.MySQLError as e:
    connection.rollback()  # 发生错误时回滚事务
    print(f"数据插入失败: {e}")

#使用DDL语句训练RAG Model
vn.train(ddl="""
    CREATE TABLE IF NOT EXISTS vuser (
        id INT PRIMARY KEY COMMENT '用户ID', 
        username VARCHAR(50) COMMENT '用户名',
        email VARCHAR(100) COMMENT '电子邮件', 
        age INT COMMENT '年龄',
        gender VARCHAR(10) COMMENT '性别（男/女）',
        city VARCHAR(50) COMMENT '城市'
    ) COMMENT='用户信息表' CHARACTER SET=utf8 COLLATE=utf8_unicode_ci;
""")

# 对话
vn.ask("年龄在30岁以上的用户有哪些?")